{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0e4d6c56",
   "metadata": {},
   "source": [
    "### Feature Importance via the attention weights\n",
    "\n",
    "I will start by saying that I consider this feature of the library purely experimental. First of all I think there are multiple ways one could address finding the features importances for these models. However, and more importantly, one has to bear in mind that even tree-based algorithms on the same dataset produce different feature importances. This is more \"dramatic\" if one uses different techniques, such as shap or feature permutation (see for example [this](https://reneelin2019.medium.com/calculating-feature-importance-with-permutation-to-explain-the-model-income-prediction-example-38a52e67441d) and references therein). All this to say that, sometimes, feature importance is just a measure contained within the experiment run, and for the model used.\n",
    "\n",
    "With that in mind, each instantiation of a deep tabular model, that has millions of trainable parameters, will potentially produce a different set of feature importances, even if the model has the same architecture. Moreover, this effect will become more apparent if the dataset is relatively easy and there are dependent/related columns so that one could get to the same success metric with different parameters. \n",
    "\n",
    "In summary, feature importances are implemented in this librray for all attention-based models for tabular data, with the exception of the `TabPerceiver`. However this functionality has to be used and interpreted with care and consider of value within the 'universe' (or context) of the model with which these features were produced.\n",
    "\n",
    "Nonetheless, let's have a look to how one would access to the feature importances when using this library. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "365df643",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "import numpy as np\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "\n",
    "from pytorch_widedeep import Trainer\n",
    "from pytorch_widedeep.models import TabTransformer, ContextAttentionMLP, WideDeep\n",
    "from pytorch_widedeep.callbacks import EarlyStopping\n",
    "from pytorch_widedeep.metrics import Accuracy\n",
    "from pytorch_widedeep.datasets import load_adult\n",
    "from pytorch_widedeep.preprocessing import TabPreprocessor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "0f935ef4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# use_cuda = torch.cuda.is_available()\n",
    "df = load_adult(as_frame=True)\n",
    "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
    "df[\"income_label\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
    "df.drop([\"income\", \"fnlwgt\", \"educational_num\"], axis=1, inplace=True)\n",
    "target_colname = \"income_label\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "202e7377",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>education</th>\n",
       "      <th>marital_status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>gender</th>\n",
       "      <th>capital_gain</th>\n",
       "      <th>capital_loss</th>\n",
       "      <th>hours_per_week</th>\n",
       "      <th>native_country</th>\n",
       "      <th>income_label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>25</td>\n",
       "      <td>Private</td>\n",
       "      <td>11th</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>Machine-op-inspct</td>\n",
       "      <td>Own-child</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>38</td>\n",
       "      <td>Private</td>\n",
       "      <td>HS-grad</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Farming-fishing</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>50</td>\n",
       "      <td>United-States</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>28</td>\n",
       "      <td>Local-gov</td>\n",
       "      <td>Assoc-acdm</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Protective-serv</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>44</td>\n",
       "      <td>Private</td>\n",
       "      <td>Some-college</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Machine-op-inspct</td>\n",
       "      <td>Husband</td>\n",
       "      <td>Black</td>\n",
       "      <td>Male</td>\n",
       "      <td>7688</td>\n",
       "      <td>0</td>\n",
       "      <td>40</td>\n",
       "      <td>United-States</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>18</td>\n",
       "      <td>?</td>\n",
       "      <td>Some-college</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>?</td>\n",
       "      <td>Own-child</td>\n",
       "      <td>White</td>\n",
       "      <td>Female</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>30</td>\n",
       "      <td>United-States</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age  workclass     education      marital_status         occupation  \\\n",
       "0   25    Private          11th       Never-married  Machine-op-inspct   \n",
       "1   38    Private       HS-grad  Married-civ-spouse    Farming-fishing   \n",
       "2   28  Local-gov    Assoc-acdm  Married-civ-spouse    Protective-serv   \n",
       "3   44    Private  Some-college  Married-civ-spouse  Machine-op-inspct   \n",
       "4   18          ?  Some-college       Never-married                  ?   \n",
       "\n",
       "  relationship   race  gender  capital_gain  capital_loss  hours_per_week  \\\n",
       "0    Own-child  Black    Male             0             0              40   \n",
       "1      Husband  White    Male             0             0              50   \n",
       "2      Husband  White    Male             0             0              40   \n",
       "3      Husband  Black    Male          7688             0              40   \n",
       "4    Own-child  White  Female             0             0              30   \n",
       "\n",
       "  native_country  income_label  \n",
       "0  United-States             0  \n",
       "1  United-States             0  \n",
       "2  United-States             1  \n",
       "3  United-States             1  \n",
       "4  United-States             0  "
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4264cba5",
   "metadata": {},
   "outputs": [],
   "source": [
    "cat_embed_cols = []\n",
    "for col in df.columns:\n",
    "    if df[col].dtype == \"O\" or df[col].nunique() < 200 and col != target_colname:\n",
    "        cat_embed_cols.append(col)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "48ae0977",
   "metadata": {},
   "outputs": [],
   "source": [
    "# all cols will be categorical\n",
    "assert len(cat_embed_cols) == df.shape[1] - 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0abe7bb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "train, test = train_test_split(\n",
    "    df, test_size=0.1, random_state=1, stratify=df[[target_colname]]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "a762aee5",
   "metadata": {},
   "outputs": [],
   "source": [
    "tab_preprocessor = TabPreprocessor(cat_embed_cols=cat_embed_cols, with_attention=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a233f9a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_tab_train = tab_preprocessor.fit_transform(train)\n",
    "X_tab_test = tab_preprocessor.transform(test)\n",
    "target = train[target_colname].values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c17233fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "tab_transformer = TabTransformer(\n",
    "    column_idx=tab_preprocessor.column_idx,\n",
    "    cat_embed_input=tab_preprocessor.cat_embed_input,\n",
    "    embed_continuous_method=\"standard\",\n",
    "    input_dim=8,\n",
    "    n_heads=2,\n",
    "    n_blocks=1,\n",
    "    attn_dropout=0.1,\n",
    "    transformer_activation=\"relu\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5010c08d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = WideDeep(deeptabular=tab_transformer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "145803d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "24133088",
   "metadata": {},
   "outputs": [],
   "source": [
    "lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
    "    optimizer,\n",
    "    threshold=0.001,\n",
    "    threshold_mode=\"abs\",\n",
    "    patience=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "06330871",
   "metadata": {},
   "outputs": [],
   "source": [
    "early_stopping = EarlyStopping(\n",
    "    min_delta=0.001, patience=30, restore_best_weights=True, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "405befbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model,\n",
    "    objective=\"binary\",\n",
    "    optimizers=optimizer,\n",
    "    lr_schedulers=lr_scheduler,\n",
    "    reducelronplateau_criterion=\"loss\",\n",
    "    callbacks=[early_stopping],\n",
    "    metrics=[Accuracy],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "997a5283",
   "metadata": {},
   "source": [
    "The feature importances will be computed after training, using a sample of the training dataset of size `feature_importance_sample_size`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "459d0f69",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1: 100%|███████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 81.80it/s, loss=0.334, metrics={'acc': 0.847}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 111.34it/s, loss=0.294, metrics={'acc': 0.8669}]\n",
      "epoch 2: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 83.02it/s, loss=0.293, metrics={'acc': 0.8656}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 124.03it/s, loss=0.283, metrics={'acc': 0.8678}]\n",
      "epoch 3: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 87.69it/s, loss=0.282, metrics={'acc': 0.8703}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 118.22it/s, loss=0.279, metrics={'acc': 0.8717}]\n",
      "epoch 4: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.24it/s, loss=0.277, metrics={'acc': 0.8718}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 115.29it/s, loss=0.277, metrics={'acc': 0.8731}]\n",
      "epoch 5: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 83.76it/s, loss=0.275, metrics={'acc': 0.8727}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 120.80it/s, loss=0.276, metrics={'acc': 0.8727}]\n",
      "epoch 6: 100%|███████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 82.78it/s, loss=0.273, metrics={'acc': 0.873}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 100.43it/s, loss=0.276, metrics={'acc': 0.871}]\n",
      "epoch 7: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.35it/s, loss=0.271, metrics={'acc': 0.8742}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 116.14it/s, loss=0.275, metrics={'acc': 0.8726}]\n",
      "epoch 8: 100%|███████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 74.29it/s, loss=0.271, metrics={'acc': 0.875}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 117.46it/s, loss=0.276, metrics={'acc': 0.8718}]\n",
      "epoch 9: 100%|███████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.10it/s, loss=0.27, metrics={'acc': 0.8761}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 105.49it/s, loss=0.275, metrics={'acc': 0.8728}]\n",
      "epoch 10: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 70.40it/s, loss=0.269, metrics={'acc': 0.8747}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 105.47it/s, loss=0.275, metrics={'acc': 0.8726}]\n",
      "epoch 11: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 72.83it/s, loss=0.268, metrics={'acc': 0.8742}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 105.03it/s, loss=0.274, metrics={'acc': 0.873}]\n",
      "epoch 12: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 71.86it/s, loss=0.267, metrics={'acc': 0.8743}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 106.61it/s, loss=0.274, metrics={'acc': 0.8734}]\n",
      "epoch 13: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 72.39it/s, loss=0.267, metrics={'acc': 0.876}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 108.05it/s, loss=0.275, metrics={'acc': 0.8717}]\n",
      "epoch 14: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 73.36it/s, loss=0.265, metrics={'acc': 0.8767}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 109.60it/s, loss=0.276, metrics={'acc': 0.8747}]\n",
      "epoch 15: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 73.34it/s, loss=0.264, metrics={'acc': 0.876}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 117.55it/s, loss=0.276, metrics={'acc': 0.8706}]\n",
      "epoch 16: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.35it/s, loss=0.264, metrics={'acc': 0.8777}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 122.08it/s, loss=0.275, metrics={'acc': 0.8753}]\n",
      "epoch 17: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.96it/s, loss=0.263, metrics={'acc': 0.877}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 120.83it/s, loss=0.277, metrics={'acc': 0.8739}]\n",
      "epoch 18: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.38it/s, loss=0.263, metrics={'acc': 0.8779}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 118.18it/s, loss=0.278, metrics={'acc': 0.8714}]\n",
      "epoch 19: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.61it/s, loss=0.261, metrics={'acc': 0.8784}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 119.19it/s, loss=0.278, metrics={'acc': 0.8712}]\n",
      "epoch 20: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 72.43it/s, loss=0.261, metrics={'acc': 0.8791}]\n",
      "valid: 100%|███████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 118.87it/s, loss=0.28, metrics={'acc': 0.873}]\n",
      "epoch 21: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 73.97it/s, loss=0.26, metrics={'acc': 0.8787}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 107.50it/s, loss=0.279, metrics={'acc': 0.8732}]\n",
      "epoch 22: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 71.76it/s, loss=0.253, metrics={'acc': 0.8816}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 108.11it/s, loss=0.279, metrics={'acc': 0.8707}]\n",
      "epoch 23: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 71.92it/s, loss=0.252, metrics={'acc': 0.8828}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 100.14it/s, loss=0.28, metrics={'acc': 0.8711}]\n",
      "epoch 24: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.04it/s, loss=0.252, metrics={'acc': 0.8829}]\n",
      "valid: 100%|███████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 98.36it/s, loss=0.28, metrics={'acc': 0.8708}]\n",
      "epoch 25: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.31it/s, loss=0.251, metrics={'acc': 0.883}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 122.97it/s, loss=0.281, metrics={'acc': 0.8709}]\n",
      "epoch 26: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.63it/s, loss=0.25, metrics={'acc': 0.8834}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 123.07it/s, loss=0.281, metrics={'acc': 0.8698}]\n",
      "epoch 27: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 78.37it/s, loss=0.251, metrics={'acc': 0.884}]\n",
      "valid: 100%|███████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 114.75it/s, loss=0.281, metrics={'acc': 0.87}]\n",
      "epoch 28: 100%|███████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.19it/s, loss=0.25, metrics={'acc': 0.883}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 117.99it/s, loss=0.282, metrics={'acc': 0.8699}]\n",
      "epoch 29: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.19it/s, loss=0.25, metrics={'acc': 0.8829}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 111.11it/s, loss=0.282, metrics={'acc': 0.8695}]\n",
      "epoch 30: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.57it/s, loss=0.249, metrics={'acc': 0.8839}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 115.64it/s, loss=0.283, metrics={'acc': 0.8689}]\n",
      "epoch 31: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.55it/s, loss=0.249, metrics={'acc': 0.8846}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 125.10it/s, loss=0.283, metrics={'acc': 0.869}]\n",
      "epoch 32: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 82.56it/s, loss=0.248, metrics={'acc': 0.8841}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 118.45it/s, loss=0.284, metrics={'acc': 0.8687}]\n",
      "epoch 33: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 81.06it/s, loss=0.248, metrics={'acc': 0.8848}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 129.26it/s, loss=0.284, metrics={'acc': 0.8689}]\n",
      "epoch 34: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 80.53it/s, loss=0.248, metrics={'acc': 0.8854}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 120.61it/s, loss=0.283, metrics={'acc': 0.869}]\n",
      "epoch 35: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 80.78it/s, loss=0.248, metrics={'acc': 0.8853}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 127.31it/s, loss=0.283, metrics={'acc': 0.8694}]\n",
      "epoch 36: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 82.51it/s, loss=0.248, metrics={'acc': 0.8863}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 125.94it/s, loss=0.283, metrics={'acc': 0.8693}]\n",
      "epoch 37: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 81.35it/s, loss=0.247, metrics={'acc': 0.8844}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 126.77it/s, loss=0.283, metrics={'acc': 0.8692}]\n",
      "epoch 38: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 80.62it/s, loss=0.248, metrics={'acc': 0.8837}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 116.62it/s, loss=0.283, metrics={'acc': 0.8692}]\n",
      "epoch 39: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.82it/s, loss=0.248, metrics={'acc': 0.8842}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 123.64it/s, loss=0.283, metrics={'acc': 0.8695}]\n",
      "epoch 40: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.86it/s, loss=0.247, metrics={'acc': 0.8855}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 120.34it/s, loss=0.283, metrics={'acc': 0.8692}]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best Epoch: 10. Best val_loss: 0.27451\n",
      "Restoring model weights from the end of the best epoch\n"
     ]
    }
   ],
   "source": [
    "trainer.fit(\n",
    "    X_tab=X_tab_train,\n",
    "    target=target,\n",
    "    val_split=0.2,\n",
    "    n_epochs=100,\n",
    "    batch_size=128,\n",
    "    validation_freq=1,\n",
    "    feature_importance_sample_size=1000,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "cd603a44",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'age': 0.09718182,\n",
       " 'workclass': 0.090637445,\n",
       " 'education': 0.08910798,\n",
       " 'marital_status': 0.08971319,\n",
       " 'occupation': 0.12546304,\n",
       " 'relationship': 0.086381145,\n",
       " 'race': 0.050686445,\n",
       " 'gender': 0.05116429,\n",
       " 'capital_gain': 0.08165918,\n",
       " 'capital_loss': 0.07702667,\n",
       " 'hours_per_week': 0.08205996,\n",
       " 'native_country': 0.07891885}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.feature_importance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8342dfdc",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "predict: 100%|████████████████████████████████████████████████████████████████████████████████████████████████| 39/39 [00:00<00:00, 199.63it/s]\n"
     ]
    }
   ],
   "source": [
    "preds = trainer.predict(X_tab=X_tab_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "7cfdabda",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8685772773797339"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy_score(preds, test.income_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5b49516b",
   "metadata": {},
   "outputs": [],
   "source": [
    "test.reset_index(drop=True, inplace=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "6ad4b3b6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>education</th>\n",
       "      <th>marital_status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>gender</th>\n",
       "      <th>capital_gain</th>\n",
       "      <th>capital_loss</th>\n",
       "      <th>hours_per_week</th>\n",
       "      <th>native_country</th>\n",
       "      <th>income_label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>26</td>\n",
       "      <td>Private</td>\n",
       "      <td>Some-college</td>\n",
       "      <td>Never-married</td>\n",
       "      <td>Exec-managerial</td>\n",
       "      <td>Not-in-family</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>0</td>\n",
       "      <td>60</td>\n",
       "      <td>United-States</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age workclass     education marital_status       occupation   relationship  \\\n",
       "0   26   Private  Some-college  Never-married  Exec-managerial  Not-in-family   \n",
       "\n",
       "    race gender  capital_gain  capital_loss  hours_per_week native_country  \\\n",
       "0  White   Male             0             0              60  United-States   \n",
       "\n",
       "   income_label  \n",
       "0             0  "
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test[test.income_label == 0].head(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "6419edff",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>age</th>\n",
       "      <th>workclass</th>\n",
       "      <th>education</th>\n",
       "      <th>marital_status</th>\n",
       "      <th>occupation</th>\n",
       "      <th>relationship</th>\n",
       "      <th>race</th>\n",
       "      <th>gender</th>\n",
       "      <th>capital_gain</th>\n",
       "      <th>capital_loss</th>\n",
       "      <th>hours_per_week</th>\n",
       "      <th>native_country</th>\n",
       "      <th>income_label</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>36</td>\n",
       "      <td>Local-gov</td>\n",
       "      <td>Doctorate</td>\n",
       "      <td>Married-civ-spouse</td>\n",
       "      <td>Prof-specialty</td>\n",
       "      <td>Husband</td>\n",
       "      <td>White</td>\n",
       "      <td>Male</td>\n",
       "      <td>0</td>\n",
       "      <td>1887</td>\n",
       "      <td>50</td>\n",
       "      <td>United-States</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   age  workclass  education      marital_status      occupation relationship  \\\n",
       "3   36  Local-gov  Doctorate  Married-civ-spouse  Prof-specialty      Husband   \n",
       "\n",
       "    race gender  capital_gain  capital_loss  hours_per_week native_country  \\\n",
       "3  White   Male             0          1887              50  United-States   \n",
       "\n",
       "   income_label  \n",
       "3             1  "
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test[test.income_label == 1].head(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6860689",
   "metadata": {},
   "source": [
    "To get the feature importance of a test dataset, simply use the `explain` method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "ba6432ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "feat_imp_per_sample = trainer.explain(X_tab_test, save_step_masks=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "f350e822",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['hours_per_week',\n",
       " 'education',\n",
       " 'relationship',\n",
       " 'occupation',\n",
       " 'workclass',\n",
       " 'capital_gain',\n",
       " 'native_country',\n",
       " 'marital_status',\n",
       " 'capital_loss',\n",
       " 'age',\n",
       " 'race',\n",
       " 'gender']"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(test.iloc[0].index[np.argsort(-feat_imp_per_sample[0])])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "f05a8302",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['age',\n",
       " 'capital_loss',\n",
       " 'hours_per_week',\n",
       " 'marital_status',\n",
       " 'native_country',\n",
       " 'relationship',\n",
       " 'race',\n",
       " 'education',\n",
       " 'occupation',\n",
       " 'capital_gain',\n",
       " 'gender',\n",
       " 'workclass']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "list(test.iloc[3].index[np.argsort(-feat_imp_per_sample[3])])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "558fc30f",
   "metadata": {},
   "source": [
    "We could do the same with the `ContextAttentionMLP`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "45b8a869",
   "metadata": {},
   "outputs": [],
   "source": [
    "context_attn_mlp = ContextAttentionMLP(\n",
    "    column_idx=tab_preprocessor.column_idx,\n",
    "    cat_embed_input=tab_preprocessor.cat_embed_input,\n",
    "    cat_embed_dropout=0.0,\n",
    "    input_dim=16,\n",
    "    attn_dropout=0.1,\n",
    "    attn_activation=\"relu\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "07f0c434",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp_model = WideDeep(deeptabular=context_attn_mlp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d9c1ef75",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp_optimizer = torch.optim.Adam(mlp_model.parameters(), lr=0.01, weight_decay=0.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "4c2f21ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp_lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
    "    mlp_optimizer,\n",
    "    threshold=0.001,\n",
    "    threshold_mode=\"abs\",\n",
    "    patience=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "504c1755",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp_early_stopping = EarlyStopping(\n",
    "    min_delta=0.001, patience=30, restore_best_weights=True, verbose=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "78c79a3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "mlp_trainer = Trainer(\n",
    "    mlp_model,\n",
    "    objective=\"binary\",\n",
    "    optimizers=mlp_optimizer,\n",
    "    lr_schedulers=mlp_lr_scheduler,\n",
    "    reducelronplateau_criterion=\"loss\",\n",
    "    callbacks=[mlp_early_stopping],\n",
    "    metrics=[Accuracy],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "af40082b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "epoch 1: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 73.11it/s, loss=0.405, metrics={'acc': 0.8094}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 119.26it/s, loss=0.309, metrics={'acc': 0.8583}]\n",
      "epoch 2: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 71.70it/s, loss=0.332, metrics={'acc': 0.8447}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 116.36it/s, loss=0.293, metrics={'acc': 0.8646}]\n",
      "epoch 3: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 78.42it/s, loss=0.319, metrics={'acc': 0.8505}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 118.05it/s, loss=0.293, metrics={'acc': 0.8654}]\n",
      "epoch 4: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.00it/s, loss=0.312, metrics={'acc': 0.8554}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 116.49it/s, loss=0.291, metrics={'acc': 0.8661}]\n",
      "epoch 5: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.50it/s, loss=0.308, metrics={'acc': 0.8583}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 120.48it/s, loss=0.287, metrics={'acc': 0.8669}]\n",
      "epoch 6: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 80.84it/s, loss=0.303, metrics={'acc': 0.8605}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 128.70it/s, loss=0.288, metrics={'acc': 0.8673}]\n",
      "epoch 7: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.93it/s, loss=0.301, metrics={'acc': 0.8597}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 121.47it/s, loss=0.298, metrics={'acc': 0.8628}]\n",
      "epoch 8: 100%|████████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 80.56it/s, loss=0.3, metrics={'acc': 0.8592}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 119.84it/s, loss=0.281, metrics={'acc': 0.8718}]\n",
      "epoch 9: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.17it/s, loss=0.298, metrics={'acc': 0.8619}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 126.32it/s, loss=0.28, metrics={'acc': 0.8716}]\n",
      "epoch 10: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 82.13it/s, loss=0.297, metrics={'acc': 0.8615}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 120.50it/s, loss=0.281, metrics={'acc': 0.8718}]\n",
      "epoch 11: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 82.54it/s, loss=0.293, metrics={'acc': 0.8641}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 122.57it/s, loss=0.284, metrics={'acc': 0.867}]\n",
      "epoch 12: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 80.92it/s, loss=0.293, metrics={'acc': 0.863}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 126.42it/s, loss=0.282, metrics={'acc': 0.8701}]\n",
      "epoch 13: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.61it/s, loss=0.293, metrics={'acc': 0.8635}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 117.56it/s, loss=0.276, metrics={'acc': 0.8719}]\n",
      "epoch 14: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 74.92it/s, loss=0.29, metrics={'acc': 0.8633}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 112.06it/s, loss=0.286, metrics={'acc': 0.8669}]\n",
      "epoch 15: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.83it/s, loss=0.291, metrics={'acc': 0.865}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 112.88it/s, loss=0.282, metrics={'acc': 0.8677}]\n",
      "epoch 16: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.74it/s, loss=0.29, metrics={'acc': 0.8653}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 120.85it/s, loss=0.285, metrics={'acc': 0.8672}]\n",
      "epoch 17: 100%|███████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.99it/s, loss=0.29, metrics={'acc': 0.865}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 113.53it/s, loss=0.282, metrics={'acc': 0.8681}]\n",
      "epoch 18: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 71.22it/s, loss=0.288, metrics={'acc': 0.8651}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 112.89it/s, loss=0.288, metrics={'acc': 0.8676}]\n",
      "epoch 19: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.85it/s, loss=0.29, metrics={'acc': 0.8661}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 114.26it/s, loss=0.284, metrics={'acc': 0.8662}]\n",
      "epoch 20: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.80it/s, loss=0.289, metrics={'acc': 0.8661}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 119.44it/s, loss=0.281, metrics={'acc': 0.8703}]\n",
      "epoch 21: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.72it/s, loss=0.29, metrics={'acc': 0.8661}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 113.04it/s, loss=0.285, metrics={'acc': 0.8648}]\n",
      "epoch 22: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 74.86it/s, loss=0.289, metrics={'acc': 0.8656}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 111.75it/s, loss=0.282, metrics={'acc': 0.8666}]\n",
      "epoch 23: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.68it/s, loss=0.289, metrics={'acc': 0.8668}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 115.86it/s, loss=0.282, metrics={'acc': 0.8724}]\n",
      "epoch 24: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.93it/s, loss=0.288, metrics={'acc': 0.8653}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 114.69it/s, loss=0.285, metrics={'acc': 0.8656}]\n",
      "epoch 25: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.99it/s, loss=0.284, metrics={'acc': 0.8671}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 118.27it/s, loss=0.277, metrics={'acc': 0.8707}]\n",
      "epoch 26: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.86it/s, loss=0.282, metrics={'acc': 0.8686}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 123.94it/s, loss=0.276, metrics={'acc': 0.8712}]\n",
      "epoch 27: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.86it/s, loss=0.283, metrics={'acc': 0.8691}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 114.11it/s, loss=0.277, metrics={'acc': 0.8716}]\n",
      "epoch 28: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.43it/s, loss=0.281, metrics={'acc': 0.8696}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 114.64it/s, loss=0.277, metrics={'acc': 0.8712}]\n",
      "epoch 29: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.39it/s, loss=0.281, metrics={'acc': 0.8696}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 115.83it/s, loss=0.277, metrics={'acc': 0.872}]\n",
      "epoch 30: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 74.93it/s, loss=0.28, metrics={'acc': 0.8706}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 113.07it/s, loss=0.275, metrics={'acc': 0.8714}]\n",
      "epoch 31: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.35it/s, loss=0.281, metrics={'acc': 0.8697}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 112.68it/s, loss=0.276, metrics={'acc': 0.872}]\n",
      "epoch 32: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.72it/s, loss=0.28, metrics={'acc': 0.8693}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 122.50it/s, loss=0.276, metrics={'acc': 0.8709}]\n",
      "epoch 33: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.68it/s, loss=0.28, metrics={'acc': 0.8716}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 110.07it/s, loss=0.277, metrics={'acc': 0.8709}]\n",
      "epoch 34: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 73.58it/s, loss=0.279, metrics={'acc': 0.8704}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 111.07it/s, loss=0.274, metrics={'acc': 0.8719}]\n",
      "epoch 35: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 74.03it/s, loss=0.28, metrics={'acc': 0.8687}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 110.15it/s, loss=0.276, metrics={'acc': 0.871}]\n",
      "epoch 36: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 73.11it/s, loss=0.279, metrics={'acc': 0.8706}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 111.10it/s, loss=0.278, metrics={'acc': 0.8705}]\n",
      "epoch 37: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 74.30it/s, loss=0.279, metrics={'acc': 0.869}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 110.75it/s, loss=0.279, metrics={'acc': 0.8702}]\n",
      "epoch 38: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 72.34it/s, loss=0.28, metrics={'acc': 0.8691}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 117.79it/s, loss=0.277, metrics={'acc': 0.8698}]\n",
      "epoch 39: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.49it/s, loss=0.279, metrics={'acc': 0.8694}]\n",
      "valid: 100%|███████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 115.25it/s, loss=0.279, metrics={'acc': 0.87}]\n",
      "epoch 40: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.51it/s, loss=0.28, metrics={'acc': 0.8694}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 110.90it/s, loss=0.277, metrics={'acc': 0.8694}]\n",
      "epoch 41: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.35it/s, loss=0.278, metrics={'acc': 0.8716}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 115.74it/s, loss=0.28, metrics={'acc': 0.8675}]\n",
      "epoch 42: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.35it/s, loss=0.279, metrics={'acc': 0.8695}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 122.76it/s, loss=0.277, metrics={'acc': 0.8699}]\n",
      "epoch 43: 100%|█████████████████████████████████████████████████████████| 275/275 [00:04<00:00, 66.14it/s, loss=0.279, metrics={'acc': 0.8681}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 106.20it/s, loss=0.277, metrics={'acc': 0.8714}]\n",
      "epoch 44: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 73.37it/s, loss=0.279, metrics={'acc': 0.8704}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 115.18it/s, loss=0.277, metrics={'acc': 0.8716}]\n",
      "epoch 45: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.23it/s, loss=0.278, metrics={'acc': 0.8702}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 124.83it/s, loss=0.278, metrics={'acc': 0.8707}]\n",
      "epoch 46: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.12it/s, loss=0.278, metrics={'acc': 0.8704}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 126.62it/s, loss=0.279, metrics={'acc': 0.8693}]\n",
      "epoch 47: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 78.55it/s, loss=0.276, metrics={'acc': 0.8713}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 117.99it/s, loss=0.279, metrics={'acc': 0.8691}]\n",
      "epoch 48: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.25it/s, loss=0.278, metrics={'acc': 0.8719}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 124.52it/s, loss=0.278, metrics={'acc': 0.8695}]\n",
      "epoch 49: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 78.35it/s, loss=0.277, metrics={'acc': 0.8721}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 119.82it/s, loss=0.279, metrics={'acc': 0.8691}]\n",
      "epoch 50: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.15it/s, loss=0.277, metrics={'acc': 0.8717}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 122.62it/s, loss=0.278, metrics={'acc': 0.8699}]\n",
      "epoch 51: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 78.55it/s, loss=0.277, metrics={'acc': 0.8713}]\n",
      "valid: 100%|███████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 117.63it/s, loss=0.278, metrics={'acc': 0.87}]\n",
      "epoch 52: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.04it/s, loss=0.276, metrics={'acc': 0.8721}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 126.39it/s, loss=0.278, metrics={'acc': 0.8697}]\n",
      "epoch 53: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 78.15it/s, loss=0.277, metrics={'acc': 0.8721}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 127.56it/s, loss=0.278, metrics={'acc': 0.8699}]\n",
      "epoch 54: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.41it/s, loss=0.277, metrics={'acc': 0.8711}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 118.95it/s, loss=0.278, metrics={'acc': 0.8698}]\n",
      "epoch 55: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 76.35it/s, loss=0.277, metrics={'acc': 0.8718}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 126.90it/s, loss=0.278, metrics={'acc': 0.8699}]\n",
      "epoch 56: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.83it/s, loss=0.277, metrics={'acc': 0.8707}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 115.13it/s, loss=0.279, metrics={'acc': 0.8691}]\n",
      "epoch 57: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.20it/s, loss=0.277, metrics={'acc': 0.8722}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 123.16it/s, loss=0.279, metrics={'acc': 0.8691}]\n",
      "epoch 58: 100%|██████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 73.33it/s, loss=0.276, metrics={'acc': 0.871}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 123.37it/s, loss=0.278, metrics={'acc': 0.8691}]\n",
      "epoch 59: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 78.41it/s, loss=0.277, metrics={'acc': 0.8714}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 125.17it/s, loss=0.278, metrics={'acc': 0.8695}]\n",
      "epoch 60: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 79.26it/s, loss=0.276, metrics={'acc': 0.8721}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 120.60it/s, loss=0.278, metrics={'acc': 0.869}]\n",
      "epoch 61: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 75.88it/s, loss=0.278, metrics={'acc': 0.8703}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 124.47it/s, loss=0.278, metrics={'acc': 0.8692}]\n",
      "epoch 62: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.13it/s, loss=0.276, metrics={'acc': 0.8711}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 125.80it/s, loss=0.278, metrics={'acc': 0.8691}]\n",
      "epoch 63: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 77.20it/s, loss=0.277, metrics={'acc': 0.8715}]\n",
      "valid: 100%|█████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 118.50it/s, loss=0.278, metrics={'acc': 0.8695}]\n",
      "epoch 64: 100%|█████████████████████████████████████████████████████████| 275/275 [00:03<00:00, 78.11it/s, loss=0.276, metrics={'acc': 0.8719}]\n",
      "valid: 100%|██████████████████████████████████████████████████████████████| 69/69 [00:00<00:00, 114.52it/s, loss=0.278, metrics={'acc': 0.869}]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Best Epoch: 34. Best val_loss: 0.27449\n",
      "Restoring model weights from the end of the best epoch\n"
     ]
    }
   ],
   "source": [
    "mlp_trainer.fit(\n",
    "    X_tab=X_tab_train,\n",
    "    target=target,\n",
    "    val_split=0.2,\n",
    "    n_epochs=100,\n",
    "    batch_size=128,\n",
    "    validation_freq=1,\n",
    "    feature_importance_sample_size=1000,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "0a4fdc77",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'age': 0.116632804,\n",
       " 'workclass': 0.050255153,\n",
       " 'education': 0.094621316,\n",
       " 'marital_status': 0.12328919,\n",
       " 'occupation': 0.107893184,\n",
       " 'relationship': 0.11747801,\n",
       " 'race': 0.054717205,\n",
       " 'gender': 0.07514235,\n",
       " 'capital_gain': 0.059732802,\n",
       " 'capital_loss': 0.06738944,\n",
       " 'hours_per_week': 0.0610674,\n",
       " 'native_country': 0.07178114}"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "mlp_trainer.feature_importance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "c10be01f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "predict: 100%|████████████████████████████████████████████████████████████████████████████████████████████████| 39/39 [00:00<00:00, 212.38it/s]\n"
     ]
    }
   ],
   "source": [
    "mlp_preds = mlp_trainer.predict(X_tab=X_tab_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "0b2a46a9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.8726714431934494"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "accuracy_score(mlp_preds, test.income_label)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
